#include "graphcut.h"

GraphCut::GraphCut(Image i, Trimap t, Model* m) :
	_oTrimap(t), _oImage(i), _oModel(m),_oGraph(i.getWidth(),i.getHeight()){
	_aProbs.resize(i.getWidth(),i.getHeight());
	VERBOSE("constructing Graphcut");
//	_oGraph = new grabCutGraph(i.getWidth(), i.getHeight());
}
GraphCut::~GraphCut() {
	MAXVERBOSE("destructing Graphcut");
//	delete _oGraph;
}
Multimap GraphCut::process() {
	VERBOSE("initialising Graph");
	//initialise Graph
	Weights w(0,0,0,0);
//	assert(_oGraph);
	for (uint x = 0; x<_oImage.getWidth(); ++x)
		for (uint y = 0; y<_oImage.getHeight(); ++y) {
			_oModel->getWeights(x,y,w,_oTrimap);
			_aProbs.setXY(x,y,int(round(w.getObject()*255)));
			_oGraph.addNode(x,y,w);
		}
	//do max-flow
	VERBOSE("maxflow");
	_oGraph.maxFlow();
	//postprocessing/alpha-blending
	VERBOSE("alhpa-blending");
	Multimap m = getSoftSegmentation(getHardSegmentation(),_oTrimap);
	return m;
}
Multimap	GraphCut::getProbs() const {
	return _aProbs;
}
Trimap GraphCut::createRibbonTrimap(const Trimap & t, const Trimap & pre, Multimap & m, int offset) const {
	VERBOSE("creating ribbon Trimap");
	int x,y,xl,yl,tmpx,tmpy;
	Trimap work;

	work.resize(t.getWidth(),t.getHeight());
	m.resize(t.getWidth(),t.getHeight());

	for (x = 0; x < int(t.getWidth()); ++x) //create working map
	for (y = 0; y < int(t.getHeight());++y) {
		Trimap::obu locale = t.getXY(x,y);
		for (xl = std::max(0,x-offset); xl<=std::min(int(t.getWidth())-1,x+offset); ++xl)
		for (yl = std::max(0,y-offset); yl<=std::min(int(t.getHeight())-1,y+offset); ++yl) { //look at a square around locale point
			tmpx = xl-x;
			tmpy = yl-y;
			if (sqrt(tmpx*tmpx + tmpy*tmpy) < offset && int(t.getXY(xl,yl))!=locale) { //look at circle around locale point
				locale = pre.getXY(x,y); //so foreset points will not be overruled
				xl = x+offset;//must not run any longer
				yl = y+offset;
			}
		}
		work.setXY(x,y,locale);
		switch (locale) {
		case Trimap::object: m.setXY(x,y,255);		break;
		case Trimap::background: m.setXY(x,y,0);	break;
		case Trimap::unknown:											break;
		}
	}
	return work;
}
Multimap GraphCut::getSoftSegmentation(const Trimap & t, const Trimap & pre) const {
	VERBOSE("creating soft segmentation");
	uint x,y;
	Multimap m;
	Trimap work = createRibbonTrimap(t,pre,m,6);
	GmdModel gmd;
//	gmd2Model gmd;
	bool valid = work.getCount(Trimap::object)>0 && work.getCount(Trimap::background)>0;
	if (valid) {
		gmd.setImage(_oImage);
		gmd.train(work);
	}
	for (x = 0; x < work.getWidth(); ++x)
	for (y = 0; y < work.getHeight(); ++y) {
		if (valid && work.getXY(x,y)==Trimap::unknown)
			m.setXY(x,y,int(round(255*gmd.getObjectProb(x,y,work))));
		else
			m.setXY(x,y,t.getXY(x,y)==Trimap::object?255:0); //so when no model can't be trained use hard segmentation
	}
	VERBOSE("done");
	return m;
}
Trimap	GraphCut::getHardSegmentation() const {
	VERBOSE("getting hard segmentation");
//	assert(_oGraph);
	Trimap m (_oImage.getWidth(), _oImage.getHeight());
	for (uint x = 0; x<_oImage.getWidth(); ++x)
	for (uint y = 0; y<_oImage.getHeight(); ++y) {
		m.setXY(x,y,_oGraph.isForeground(x,y)?Trimap::object:Trimap::background); //binary segmentation first
	}
	VERBOSE("done");
	return m;
}
const Image			GraphCut::getImage()	const			{	return _oImage; }
const Model*		GraphCut::getModel()	const			{	return _oModel; }
const Trimap		GraphCut::getTrimap()	const			{	return _oTrimap; }
